package practise.leetcode.q4;

import java.util.ArrayList;
import java.util.List;

public class Solution {
    public double findMedianSortedArrays(int[] nums1, int[] nums2) {
        int length = nums1.length + nums2.length;
        if (length % 2 == 1) {
            int idx = length / 2;
            List<Integer> results = findNum(idx, nums1, nums2);
            if (results.isEmpty()) {
                return 0;
            }
            return results.get(0);
        }
        else {
            int lastIdx = length / 2;
            int firstIdx = lastIdx - 1;
            List<Integer> results = findNum(firstIdx, nums1, nums2);
            if (results.isEmpty()) {
                return 0;
            }
            double a = results.get(0);
            double b = results.get(1);
            return (a + b) / 2;
        }
    }

    public List<Integer> findNum(int i, int[] nums1, int[] nums2) {
        int leftIdx = 0;
        int rightIdx = 0;
        int sum = 0;
        List<Integer> results = new ArrayList<>(2);
        while ((leftIdx < nums1.length && rightIdx < nums2.length) && sum < i + 2) {
            int left = nums1[leftIdx];
            int right = nums2[rightIdx];
            boolean isMatched = sum == i || sum == i + 1;
            if (left <= right) {
                leftIdx++;
                if (isMatched) {
                    results.add(left);
                }
            }
            else {
                rightIdx++;
                if (isMatched) {
                    results.add(right);
                }
            }

            sum++;
        }

        while (leftIdx < nums1.length && sum < i + 2) {
            boolean isMatched = sum == i || sum == i + 1;
            if (isMatched) {
                results.add(nums1[leftIdx++]);
                sum++;
            }
        }

        while (rightIdx < nums2.length && sum < i + 2) {
            boolean isMatched = sum == i || sum == i + 1;
            if (isMatched) {
                results.add(nums2[rightIdx++]);
                sum++;
            }
        }

        return results;
    }
}
